from LSTM import predict_stock as ps
from sysdata import sys_base_data as sysBaseData
from LSTM.variableset import EMVar,EMPath

def predict(code,ktype='D', tag='DEFAULT'):
    from LSTM.variableset import  EMPath
    try:
        sfund = sysBaseData.get_self_fund(code)
        df = sfund.predict_data_source(to_csv=False, ktype=ktype)['data']
        df_test_origin, test_data_source, test_total_size = ps.prepare_test_data(df)
        checkpoint_dir = EMPath.check_point_dir(code=code, is_common=True, ktype=ktype)
        buyModel, _ = ps.prediction(df_test=df_test_origin,
                                    orient_data_test=test_data_source,
                                    checkpoint_dir=checkpoint_dir,
                                    test_code=code, to_csv=False, to_lean=False)
        # buyModel.save_to_mongo(tag=tag)
    except Exception as e:
        print('%s %s 模型 预测失败'%(sfund.code_with_name(), ktype))
        print(e)
    finally:
        ps.reuse_scope = True


def train(code,ktype='D', tag='DEFAULT'):
    try:
        sfund = sysBaseData.get_self_fund(code)
        df = sfund.predict_data_source(to_csv=False, ktype=ktype)['data']
        df_test_origin, test_data_source, test_total_size = ps.prepare_test_data(df)
        ps.train_lstm(test_data_source, ktype=ktype)

        # buyModel.save_to_mongo(tag=tag)
    except Exception as e:
        print('%s %s 模型 预测失败' % (sfund.code_with_name(), ktype))
        print(e)
    finally:
        pass


def predict_indicator(code,ktype='D', tag='DEFAULT'):
    from LSTM.variableset import  EMPath
    try:
        sfund = sysBaseData.get_self_fund(code)
        df = sfund.get_index(start='2010-01-01', ktype=ktype, insertdb=False, append_now=True,
                             indicator=['macd', 'max', 'adosc'])
        dp5_max = df[EMVar.max].shift(-5)
        df = df['2016-01-01':]
        dp5_max2close = (dp5_max - df[EMVar.close]) / df[EMVar.close]
        # df[EMVar.label] = dp5_max2close[dp5_max2close > 0.1]
        df[EMVar.label] = dp5_max2close
        df = df.dropna()
        df = df[[EMVar.code, EMVar.date, EMVar.macd, EMVar.macdsignal, EMVar.macdhist, EMVar.adosc, EMVar.label]]
        ps.input_size = 4
        df_test_origin, test_data_source, test_total_size = ps.prepare_test_data(df)
        checkpoint_dir = EMPath.check_point_dir(code=code, is_common=True, ktype=ktype)
        ps.prediction(df_test=df_test_origin,
                      orient_data_test=test_data_source,
                      checkpoint_dir=checkpoint_dir,
                      test_code=code, to_csv=True, to_lean=False)
        # buyModel.save_to_mongo(tag=tag)
    except Exception as e:
        print('%s %s 模型 预测失败'%(sfund.code_with_name(), ktype))
        print(e)
    finally:
        ps.reuse_scope = True

def train_indicator(code,ktype='D', tag='DEFAULT'):

    try:
        sfund = sysBaseData.get_self_fund(code)
        df = sfund.get_index(start='2010-01-01', ktype=ktype, insertdb=False, append_now=True,indicator=['macd', 'max', 'adosc'])
        dp5_max = df[EMVar.max].shift(-5)
        df = df['2016-01-01':]
        dp5_max2close = (dp5_max - df[EMVar.close])/df[EMVar.close]
        # df[EMVar.label] = dp5_max2close[dp5_max2close > 0.1]
        df[EMVar.label] = dp5_max2close
        df = df.dropna()
        df = df[[EMVar.code, EMVar.date, EMVar.macd, EMVar.macdsignal, EMVar.macdhist, EMVar.adosc,EMVar.label]]
        ps.input_size = 4

        df.to_csv(path_or_buf=EMPath.data_file_full_path(code=code, fileName='indicator_%s.csv'%code))

        df_test_origin, test_data_source, test_total_size = ps.prepare_test_data(df)
        ps.train_lstm(test_data_source, ktype=ktype)

    except Exception as e:
        print('%s %s 模型 预测失败' % (sfund.code_with_name(), ktype))
        print(e)
    finally:
        pass


if __name__ == '__main__':
    # train_indicator('002056')
    # predict_indicator('002339', ktype='D')
    pass